package world.snowcrystal.nova.controller;

import cn.hutool.json.JSONUtil;
import com.zhipu.oapi.service.v4.model.ModelData;
import io.reactivex.Flowable;
import io.reactivex.Scheduler;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import world.snowcrystal.nova.ai.AITemplate;
import world.snowcrystal.nova.ai.command.AiGenerationQuestionCommand;
import world.snowcrystal.nova.common.BaseResponse;
import world.snowcrystal.nova.common.ErrorCode;
import world.snowcrystal.nova.common.ResultUtils;
import world.snowcrystal.nova.exception.ThrowUtils;
import world.snowcrystal.nova.model.dto.question.QuestionDto;
import world.snowcrystal.nova.model.entity.App;
import world.snowcrystal.nova.model.enums.AppScoringTypeEnum;
import world.snowcrystal.nova.service.AppService;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
@RestController
@RequestMapping
public class AiController {

    private static final String DEFAULT_SYSTEM_MESSAGE = """
            你是一位严谨的出题专家，我会给你如下信息：
            ```
            应用名称，
            【【【应用描述】】】，
            应用类别，
            要生成的题目数，
            每个题目的选项数
            ```

            请你根据上述信息，按照以下步骤来出题：
            1. 要求：题目和选项尽可能地短，题目不要包含序号，每题的选项数以我提供的为主，题目不能重复
            2. 严格按照下面的 json 格式输出题目和选项
            ```
            [{"options":[{"value":"选项内容","key":"A"},{"value":"","key":"B"}],"title":"题目标题"}]
            ```
            title 是题目，options 是选项，每个选项的 key 按照英文字母序（比如 A、B、C、D）以此类推，value 是选项内容
            3. 检查题目是否包含序号，若包含序号则去除序号
            4. 返回的题目列表格式必须为 JSON 数组""";
    @Resource
    private AppService appService;

    @Resource
    private AITemplate aiTemplate;


    @Resource
    private Scheduler scheduler;

    @PostMapping("/ai/post")
    public BaseResponse<List<QuestionDto>> generateQuestion(
            @RequestBody AiGenerationQuestionCommand command) {
        // validate param

        ThrowUtils.throwIf(command == null, ErrorCode.PARAMS_ERROR);

        // 获取 command 中的参数
        Integer questionNumber = command.getQuestionNumber();
        Integer optionNumber = command.getOptionNumber();
        Long appId = command.getAppId();

        App app = appService.getById(appId);
        ThrowUtils.throwIf(app == null, ErrorCode.NOT_FOUND_ERROR);
        //MBTI 性格测试，
        //【【【快来测测你的 MBTI 性格】】】，
        //测评类，
        //10，
        //3
        String prompt = String.format("%s,\n【【【%s】】】,\n%s,\n%d,\n%d",
                app.getAppName(),
                app.getAppDesc(),
                AppScoringTypeEnum.get(app.getAppType()).getText(),
                questionNumber,
                optionNumber);

        // 封装 Prompt
        String result = aiTemplate.chat(DEFAULT_SYSTEM_MESSAGE, prompt, false, 0.75f);
        int start = result.indexOf('[');
        int end = result.lastIndexOf(']');
        result = result.substring(start, end + 1);
        List<QuestionDto> questionDtos = JSONUtil.toList(result, QuestionDto.class);
        return ResultUtils.success(questionDtos);
    }

    @GetMapping("/ai/post/stream")
    public SseEmitter generateQuestionStream(AiGenerationQuestionCommand command) {
        // validate param

        ThrowUtils.throwIf(command == null, ErrorCode.PARAMS_ERROR);


        SseEmitter sseEmitter = new SseEmitter();

        // 获取 command 中的参数
        Integer questionNumber = command.getQuestionNumber();
        Integer optionNumber = command.getOptionNumber();
        Long appId = command.getAppId();

        App app = appService.getById(appId);
        ThrowUtils.throwIf(app == null, ErrorCode.NOT_FOUND_ERROR);
        //MBTI 性格测试，
        //【【【快来测测你的 MBTI 性格】】】，
        //测评类，
        //10，
        //3
        String prompt = String.format("%s,\n【【【%s】】】,\n%s,\n%d,\n%d",
                app.getAppName(),
                app.getAppDesc(),
                AppScoringTypeEnum.get(app.getAppType()).getText(),
                questionNumber,
                optionNumber);

        // 封装 Prompt
        Flowable<ModelData> modelDataFlowable = aiTemplate.chatStream(DEFAULT_SYSTEM_MESSAGE, prompt, 0.75f);
        AtomicInteger counter = new AtomicInteger(0);
        StringBuffer buffer = new StringBuffer();

        // elect scheduler

        //TODO 普通用户使用固定数量线程的线程池，VIP 用户使用多线程线程池
//        Scheduler scheduler ;
//        if(true){
//            scheduler = this.scheduler;
//        }else{
//            scheduler = Schedulers.io();
//        }


        modelDataFlowable.observeOn(scheduler)

                .map(modelData -> modelData.getChoices().get(0)
                        .getDelta()
                        .getContent()
                )
                .map(stringPart -> stringPart.replace("\\s", ""))
                .filter(StringUtils::hasText)
                .flatMap(stringPart -> {
                    char[] chars = stringPart.toCharArray();
                    List<Character> charList = new ArrayList<>(chars.length);
                    for (char c : chars) {
                        charList.add(c);
                    }
                    return Flowable.fromIterable(charList);
                })
                .doOnNext(c -> {
                    if (c == '{') {
                        counter.incrementAndGet();
                    }
                    // 只有遇到首个 { 时，才开始 append
                    if (counter.get() > 0) {
                        buffer.append(c);
                    }
                    if (c == '}') {
                        counter.decrementAndGet();
                        if (counter.get() == 0) {
                            sseEmitter.send(buffer.toString());
                            buffer.setLength(0);
                        }
                    }
                })
                .doOnError(sseEmitter::completeWithError)
                .doOnComplete(sseEmitter::complete)
                .subscribe();
        return sseEmitter;
    }
}
